import os
from src.ai.langchain.init_llm import get_llm
from operator import itemgetter

import bs4
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_chroma import Chroma

from langchain_community.document_loaders import WebBaseLoader
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.tools import TavilySearchResults, QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.prebuilt import chat_agent_executor

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

os.environ["AZURE_OPENAI_ENDPOINT"] = "http://menshen.test.xdf.cn"
# os.environ["OPENAI_API_BASE"] = "http://menshen.test.xdf.cn"
os.environ["OPENAI_API_KEY"] = "c8575027653b42b1b47747f0b4ab135b"
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"

llm = get_llm()

USER_NAME = 'ourcrm_admin'
PWD = 'ourcrm_admin'
HOST = '172.24.30.115'
PORT = 3306
DATABASE = 'scrm_test'

# 使用mysqlclient驱动
mysql_uri = f'mysql+mysqldb://{USER_NAME}:{PWD}@{HOST}:{PORT}/{DATABASE}?charset=utf8mb4'

# 使用include_tables的原因是数据库中的表太多了，大模型会把所有的表名或者全库的数据结构都加入到上下文进行分析，所以会造成token数据超出的限制，所以只选择一部分表进行测试
# 具体错误信息：openai.BadRequestError: Error code: 400 - {'error': {'message': "This model's maximum context length is 128000 tokens. However, your messages resulted in 588933 tokens. Please reduce the length of the messages.", 'type': 'invalid_request_error', 'param': 'messages', 'code': 'context_length_exceeded'}}
mysql_db = SQLDatabase.from_uri(
    mysql_uri,
    include_tables=["crm_sys_user"]
)
# print(mysql_db.run("SELECT * FROM crm_sys_user order by id asc limit 20;"))

sql_prompt = PromptTemplate.from_template("请根据下面的问题，生成一个查询sql语句,返回结果不要使用任何 Markdown 格式，只用纯文本格式输出：{input}")

# 这个方法只会返回查询的sql语音，不会执行查询
create_sql_chain = create_sql_query_chain(llm=llm, db=mysql_db,prompt=sql_prompt) | StrOutputParser()
# resp = create_sql_chain.invoke({"question": "请问：crm_sys_user表中有多少条数据？"})
# print(resp)

print('分析结果：', create_sql_chain.invoke({"question": "请问：crm_sys_user表中有多少条数据？"}))

answer_prompt_msg = """
请根据用户提出的问题，SQL语句和SQL执行后的结果回答用户问题
    问题：{question}
    SQL查询：{query}
    SQL执行结果: {result}
"""

answer_prompt = PromptTemplate.from_template(answer_prompt_msg)

# 创建一个执行sql语句的工具
execute_sql_tool = QuerySQLDataBaseTool(db=mysql_db)

# 创建一个调用链，1、生成SQL 2.执行sql
# 模板
chain = (RunnablePassthrough.assign(query=create_sql_chain | StrOutputParser()).assign(result=itemgetter("query") | execute_sql_tool)
         | answer_prompt
         | llm
         | StrOutputParser()
         )

resp = chain.invoke({"question": "请问：crm_sys_user表中有多少条数据？"})
print(resp)